# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/8 14:02 
@Author  : ZhangShenao 
@File    : 1.使用LLMChain简化链调用.py 
@Desc    : 使用LLMChain简化链调用
"""
import dotenv
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_community.chat_models import ChatZhipuAI
from langchain_core.documents import Document
from langchain_core.prompts import ChatPromptTemplate

# 系统提示
SYSTEM_PROMPT = '''你是一个功能强大的聊天机器人，请根据之前的历史聊天上下文，回答用户的问题。
<历史上下文开始>
{context}
<历史上下文结束>
'''

# 加载环境变量
dotenv.load_dotenv()

# 编排Prompt
prompt = ChatPromptTemplate.from_messages([
    ('system', SYSTEM_PROMPT),
    ('human', '{query}'),
])

# 构建LLM
llm = ChatZhipuAI(model='glm-4-air')

# 通过函数调用,创建内置的chain
chain = create_stuff_documents_chain(prompt=prompt, llm=llm)

# 构建文档上下文
documents = [
    Document(page_content='小明喜欢绿色，但不喜欢黄色'),
    Document(page_content='小王喜欢粉色，也有一点喜欢红色'),
    Document(page_content='小泽喜欢蓝色，但更喜欢青色')
]

# 调用chain
result = chain.invoke({
    'query': '大家都喜欢什么颜色？',
    'context': documents,
})
print(result)
